#include "polyarith.h"
#include "karatsuba.h"
#include <cstdlib>

struct uint64_mod_t {
    const uint64_t &modulus;
    uint64_t val;
};

const uint64_mod_t &operator -= (uint64_mod_t &lhs, const uint64_mod_t &rhs)
{
    lhs.val = SubUintMod(lhs.val, rhs.val, lhs.modulus);
    return lhs;
}

const uint64_mod_t &operator += (uint64_mod_t &lhs, const uint64_mod_t &rhs)
{
    lhs.val = AddUintMod(lhs.val, rhs.val, lhs.modulus);
    return lhs;
}

uint64_mod_t operator*(const uint64_mod_t &lhs, const uint64_mod_t &rhs)
{
    uint64_mod_t result = { lhs.modulus, 0 };
    result.val = MultiplyUintMod(lhs.val, rhs.val, lhs.modulus);
    return result;
}

/*
 * @brief 多项式加法，考虑缩放因子
 * 将 `rhs` 中的每个项乘以 `scale` 并进行模运算，然后加到 `lhs` 中相应的项上。
 * @param lhs 输入输出参数，表示左侧多项式。它将更新为 `lhs + scale * rhs`。
 * @param rhs 输入参数，表示右侧多项式。
 * @param mod 模数，所有计算将在 `mod` 模下进行。
 * @param scale 缩放因子，`rhs` 中的每一项都将乘以此因子。
 * @param powerXScale 幂次扩展，`rhs` 中的项将偏移 `powerXScale` 个位置加到 `lhs` 上。
 */
void PolyAdd(poly &lhs, const poly &rhs, const uint64_t &mod, uint64_t scale, size_t powerXScale)
{
    lhs.resize(std::max(lhs.size(), rhs.size() + powerXScale));

    for (size_t i = 0; i < rhs.size(); ++i) {
        lhs[i + powerXScale] = AddUintMod(lhs[i + powerXScale], MultiplyUintMod(rhs[i], scale, mod), mod);
    }
}
/*
 * @brief 对多项式进行缩放
 * 该函数将多项式 `p` 中的每个项乘以指定的 `scale` 值，并在给定的模数 `mod` 下进行模运算。
 * @param p 输入输出参数，表示需要缩放的多项式。该多项式将被修改，每个项都乘以 `scale`。
 * @param scale 缩放因子，每个多项式项都将乘以该因子。
 * @param mod 模数，所有的乘法操作将在此模数下进行模运算。
 */
void PolyScale(poly &p, uint64_t scale, const uint64_t &mod)
{
    for (unsigned long &i : p) {
        MulMod(i, i, scale, mod);
    }
}
/*
 * @brief 对两个多项式进行模乘法运算
 * 该函数实现了两个多项式 `lhs` 和 `rhs` 的乘法，并对结果进行模运算。使用了 Karatsuba 算法来提高多项式乘法的效率。
 * 最终的结果会根据给定的模数 `mod` 和多项式模数 `pmod` 进行约简。
 * @param lhs 第一个输入多项式，乘法的左操作数。
 * @param rhs 第二个输入多项式，乘法的右操作数。
 * @param mod 运算的模数，所有的加法、减法和乘法都会在此模数下进行。
 * @param pmod 多项式模数，指定了对乘积结果进行的模约简。
 * @return poly 返回经过模运算后的乘积多项式，大小为 `lhs.size() + rhs.size()`。
 */
poly PolyMulMod(const poly &lhs, const poly &rhs, const uint64_t &mod, const PolyModulus &pmod)
{
    if (lhs.size() == 0 || rhs.size() == 0) {
        poly result;
        result.resize(pmod.n);
        return result;
    }

    poly result;
    result.resize(lhs.size() + rhs.size());
    assert(lhs.size() <= pmod.n);
    assert(rhs.size() <= pmod.n);
    karatsuba::Karatsuba<4>(
        &result[0], result.size(), &lhs[0], lhs.size(), &rhs[0], rhs.size(),
        [&mod](uint64_t a, uint64_t b) { return AddUintMod(a, b, mod); },
        [&mod](uint64_t a, uint64_t b) { return SubUintMod(a, b, mod); },
        [&mod](uint64_t a, uint64_t b) { return MultiplyUintMod(a, b, mod); });

    PolyReduceMod(result, mod, pmod);

    return result;
}
/*
 * @brief 对多项式进行模约简
 * 该函数将输入多项式 `lhs` 通过多项式模数 `pmod` 进行约简，使得结果的大小不超过 `pmod.n`。
 * 如果多项式的大小大于或等于 `pmod.n`，会通过减去多项式 `pmod.xPowerN` 的适当倍数来使多项式符合模数约简。
 * @param lhs 输入并输出的多项式。该多项式会在函数执行后被修改为经过模约简后的版本。
 * @param mod 运算的模数，所有的加法操作都会在此模数下进行。
 * @param pmod 多项式模数，定义了如何进行多项式的模约简。
 */
void PolyReduceMod(poly &lhs, const uint64_t &mod, const PolyModulus &pmod)
{
    if (lhs.size() < pmod.n) {
        return;
    }
    for (int64_t i = lhs.size() - 1; i >= static_cast<int64_t>(pmod.n); --i) {
        PolyAdd(lhs, pmod.xPowerN, mod, lhs[i], i - pmod.n);
    }
    lhs.resize(pmod.n);
}
/*
 * @brief 对多项式进行模幂运算
 * 该函数实现了多项式的模幂运算，即计算给定多项式 `basis` 的 `exp` 次幂（模 `mod`），结果会根据给定的多项式模数 `pmod`
 * 进行模约简。
 * @param basis 输入的基础多项式，即幂运算的底数。
 * @param exp 幂次，表示需要对 `basis` 多项式进行多少次乘法运算。
 * @param mod 运算的模数，所有的加法、乘法操作都会在此模数下进行。
 * @param pmod 多项式模数，定义了如何进行多项式的模约简，尤其是模约简时的多项式操作。
 * @return 返回经过模约简的多项式，即 `basis` 的 `exp` 次幂（模 `mod`）。
 */
poly PolyExponentiateMod(const poly &basis, size_t exp, const uint64_t &mod, const PolyModulus &pmod)
{
    bool changed = false;
    poly result = { 1 };
    for (int64_t i = std::numeric_limits<size_t>().digits - 1; i >= 0; --i) {
        if (((exp >> i) & 1) == 1) {
            result = PolyMulMod(basis, PolyMulMod(result, result, mod, pmod), mod, pmod);
            changed = true;
        } else if (changed) { // a performance optimization, no sense in squaring 1
            result = PolyMulMod(result, result, mod, pmod);
        }
    }
    return result;
}
/*
 * @brief 执行多项式的减法乘法操作
 * 该函数实现了对两个多项式 `lhs` 和 `rhs` 的减法乘法运算。
 * 该操作为 `dst[i + j] -= lhs[i] * rhs[j]`。
 * @param dst 存储结果的目标多项式，最终计算结果会存放在这个多项式中。
 * @param lhs 第一个输入多项式，参与乘法和减法操作。
 * @param rhs 第二个输入多项式，参与乘法和减法操作。
 * @param mod 运算的模数，所有的加法、减法和乘法操作都将在此模数下进行。
 */
void PolySubMul(poly &dst, const poly &lhs, const poly &rhs, const uint64_t &mod)
{
    if (lhs.size() == 0 || rhs.size() == 0) {
        return;
    }
    dst.resize(std::max(dst.size(), lhs.size() + rhs.size() - 1));
    for (size_t i = 0; i < lhs.size(); ++i) {
        for (size_t j = 0; j < rhs.size(); ++j) {
            dst[i + j] = SubUintMod(dst[i + j], MultiplyUintMod(lhs[i], rhs[j], mod), mod);
        }
    }
}
/*
 * @brief 执行多项式的加法乘法操作
 * 该函数实现了对两个多项式 `lhs` 和 `rhs` 的加法乘法运算。
 * 该操作为 `dst[i + j] += lhs[i] * rhs[j]`。
 * @param dst 存储结果的目标多项式，最终计算结果会存放在这个多项式中。
 * @param lhs 第一个输入多项式，参与乘法和加法操作。
 * @param rhs 第二个输入多项式，参与乘法和加法操作。
 * @param mod 运算的模数，所有的加法和乘法操作都将在此模数下进行。
 */
void PolyAddMul(poly &dst, const poly &lhs, const poly &rhs, const uint64_t &mod)
{
    if (lhs.size() == 0 || rhs.size() == 0) {
        return;
    }
    dst.resize(std::max(dst.size(), lhs.size() + rhs.size() - 1));
    for (size_t i = 0; i < lhs.size(); ++i) {
        for (size_t j = 0; j < rhs.size(); ++j) {
            dst[i + j] = AddUintMod(dst[i + j], MultiplyUintMod(lhs[i], rhs[j], mod), mod);
        }
    }
}
/*
 * @brief 扩展欧几里得算法（Extended Euclidean Algorithm）
 * 该函数使用扩展欧几里得算法计算给定整数 `a` 和 `b` 的最大公约数（GCD），并且返回一对系数 `(x, y)`
 * 同时计算出 `x` 和 `y`。
 * @param a 第一个整数
 * @param b 第二个整数
 * @return 返回一个包含系数 `(x, y)` 的元组，表示线性组合的系数。
 */
std::tuple<int64_t, int64_t> Eea(uint64_t a, uint64_t b)
{
    int64_t sa = 1;
    int64_t ta = 0;
    int64_t sb = 0;
    int64_t tb = 1;

    while (b != 0) {
        uint64_t quo = a / b;
        uint64_t rem = a % b;
        ta = ta - quo * tb;
        sa = sa - quo * sb;
        a = rem;
        std::swap(a, b);
        std::swap(sa, sb);
        std::swap(ta, tb);
    }
    return std::make_tuple(sa, ta);
}
/*
 * @brief 计算多项式在给定值 `x` 上的值
 * 该函数用于计算多项式 `p` 在 `x` 上的值，所有运算均在给定的模数 `mod` 下进行。
 * 多项式 `p` 的形式为：
 * @param p 多项式，表示为一个整数向量 `poly`。
 * @param x 输入值，表示在多项式 `p` 中求值的 `x`。
 * @param mod 模数，所有运算都将在该模数下进行。
 * @return 返回多项式 `p` 在 `x` 处的值，计算结果会被模 `mod`。
 */
uint64_t PolyEval(const poly &p, uint64_t x, const uint64_t &mod)
{
    uint64_t result = p[p.size() - 1];
    for (int64_t i = p.size() - 2; i >= 0; --i) {
        result = MultiplyUintMod(result, x, mod);
        result = AddUintMod(result, p[i], mod);
    }
    return result;
}
/*
 * @brief 多项式除法（带余数）
 * 该函数实现了多项式的除法运算，并返回商项。余数存储在 `lhs` 中。并将商和余数进行归一化。
 * @param lhs 输入的多项式，被除数，经过除法运算后会包含余数。
 * @param rhs 除数，必须是一个非零的多项式，且其最高次项的系数必须为1。
 * @param mod 模数，所有运算均在此模数下进行。
 *
 * @return 返回商多项式，商的次数为 `lhs.size() - rhs.size() + 1`。
 */
poly PolyDivRem(poly &lhs, const poly &rhs, const uint64_t &mod)
{
    assert(rhs.size() > 0);
    assert(rhs[rhs.size() - 1] == 1);
    if (rhs.size() > lhs.size()) {
        return {};
    }
    poly result;
    result.resize(lhs.size() - rhs.size() + 1);
    for (int64_t i = lhs.size() - 1; i >= static_cast<int64_t>(rhs.size() - 1); --i) {
        result[i - rhs.size() + 1] = lhs[i];
        PolySub(lhs, rhs, mod, lhs[i], i - rhs.size() + 1);
    }
    PolyNormalize(result);
    PolyNormalize(lhs);
    return result;
}
/*
 * @brief 扩展欧几里得算法对多项式的实现
 *
 * 该函数实现了多项式的扩展欧几里得算法
 * 其中 `sa` 和 `ta` 是贝祖等式的系数，`lhs` 和 `rhs` 是输入的多项式，返回的 `lhs` 是它们的最大公约数（gcd）。
 * 具体的算法步骤包括：
 * 1. 对输入多项式进行规范化，确保每一步操作的系数在模 `mod` 下是有效的。
 * 2. 使用扩展欧几里得算法的步骤，通过多次求余和更新来逐步计算 `sa` 和 `ta`。
 * 3. 返回的结果包含：`sa`、`ta` 和 `lhs`（最终的最大公约数）。
 * @param lhs 输入的第一个多项式。
 * @param rhs 输入的第二个多项式。
 * @param mod 模数，所有运算均在此模数下进行。
 * @return 返回一个三元组：
 * - `sa`：系数 `sa`。
 * - `ta`：系数 `ta`。
 * - `lhs`：返回 `lhs` 和 `rhs` 的最大公约数。
 */
std::tuple<poly, poly, poly> PolyEea(poly lhs, poly rhs, const uint64_t &mod)
{
    PolyNormalize(lhs);
    PolyNormalize(rhs);
    poly sa = { 1 };
    poly ta = {};
    poly sb = {};
    poly tb = { 1 };

    while (rhs.size() > 0) {
        if (rhs[rhs.size() - 1] != 1) {
            int64_t lc_inv = InvMod_(rhs[rhs.size() - 1], mod);
            PolyScale(rhs, lc_inv, mod);
            PolyScale(sb, lc_inv, mod);
            PolyScale(tb, lc_inv, mod);
        }
        poly quo = PolyDivRem(lhs, rhs, mod);
        PolySubMul(ta, quo, tb, mod);
        PolySubMul(sa, quo, sb, mod);
        std::swap(lhs, rhs);
        std::swap(sa, sb);
        std::swap(ta, tb);
    }
    return std::make_tuple(sa, ta, lhs);
}
/*
 * @brief 计算整数的模逆元
 * 该函数使用扩展欧几里得算法计算一个数 `x` 在模 `mod` 下的逆元：
 * @param x 输入值，要求与 `mod` 互质（即 `gcd(x, mod) = 1`）。
 * @param mod 模数，用于模运算。
 * @return 返回 `x` 在模 `mod` 下的逆元，如果 `x` 和 `mod` 互质且逆元存在。
 */
uint64_t InvMod_(uint64_t x, const uint64_t &mod)
{
    int64_t result = std::get<0>(Eea(x, mod));
    if (result < 0) {
        result += mod;
    }
    assert(result >= 0);
    assert(static_cast<uint64_t>(result) < mod);
    assert(MultiplyUintMod(x, result, mod) == 1);
    return static_cast<uint64_t>(result);
}

bool IsZero(const poly &x)
{
    for (uint64_t a : x) {
        if (a != 0) {
            return false;
        }
    }
    return true;
}
/*
 * @brief 判断多项式是否不可约
 *
 * 该函数检查给定的多项式 `x` 是否在模 `primeModulus` 下不可约。
 * 多项式不可约的定义是，不能被低次多项式整除。即，`x` 不能分解为两个低次多项式的乘积。
 *
 * 该算法基于多项式的不可约性判定方法，包括使用扩展模指数、模约化等技术。
 *
 * @param x 输入多项式，表示为 `std::vector<uint64_t>`
 * @param primeModulus 一个素数模数，用于模运算。
 *
 * @return 如果多项式 `x` 是不可约的，则返回 `true`，否则返回 `false`。
 */
bool IsIrreducible(const poly &x, const uint64_t &primeModulus)
{
    assert(x.size() > 0);
    if (x.size() <= 2) {
        return true;
    }
    assert(x[x.size() - 1] == 1);
    PolyModulus modX;
    modX.n = x.size() - 1;
    poly modPoly = x;
    modPoly[x.size() - 1] = 0;
    PolyNormalize(modPoly);
    PolyScale(modPoly, NegateUintMod(1, primeModulus), primeModulus);
    modX.xPowerN = std::move(modPoly);

    const size_t p = primeModulus;
    const size_t d = x.size() - 1;
    poly tmp = PolyExponentiateMod({ 0, 1 }, IPow(p, d), primeModulus, modX);
    PolySub(tmp, { 0, 1 }, primeModulus);
    PolyReduceMod(tmp, primeModulus, modX);
    PolyNormalize(tmp);
    if (tmp.size() > 0) {
        return false;
    }

    std::vector<uint64_t> primes;
    for (uint64_t n = 2; n <= d; ++n) {
        if (p == 0) {
            cout << "p==0!!!" << endl;
        }
        if (std::any_of(primes.cbegin(), primes.cend(), [n](uint64_t p) { return n % p == 0; })) {
            continue;
        }
        primes.push_back(n);

        poly tmpApp = PolyExponentiateMod({ 0, 1 }, IPow(p, d / n), primeModulus, modX);
        PolySub(tmpApp, { 0, 1 }, primeModulus);
        PolyReduceMod(tmpApp, primeModulus, modX);
        PolyNormalize(tmpApp);
        if (tmpApp.size() == 0) {
            return false;
        }
    }

    return true;
}
/*
 * @brief 多项式算子基本功能测试
 *
 */
void TestIsIrreducible()
{
    {
        uint64_t mod(2);
        assert(IsIrreducible({ 0, 1 }, mod));
        assert(IsIrreducible({ 1, 1, 1 }, mod));
        assert(!IsIrreducible({ 1, 0, 1 }, mod));
        assert(!IsIrreducible({ 0, 0, 1 }, mod));
        assert(IsIrreducible({ 1, 0, 1, 1 }, mod));
    }
    {
        uint64_t mod(127);
        assert(IsIrreducible({ 115, 22, 36, 1 }, mod));
        assert(!IsIrreducible({ 1, 0, 0, 1 }, mod));
        assert(!IsIrreducible({ 1, 1, 0, 1 }, mod));
    }
    std::cout << "TestIsIrreducible(): success" << std::endl;
}